use log;
use pnet::packet::udp::MutableUdpPacket;
use pnet_packet::{ip::IpNextHeaderProtocols, ipv4::MutableIpv4Packet, udp::ipv4_checksum, Packet};
use socket2::{Domain, Protocol, Socket, Type};
use std::net::{Ipv4Addr, SocketAddrV4};

pub struct RawUdpSocket {
    sock: Socket,
}

impl RawUdpSocket {
    pub fn new() -> Self {
        let sock = Socket::new(Domain::IPV4, Type::RAW, Some(Protocol::UDP))
            .expect("create socket failed");
        Self { sock }
    }
    pub fn probe_local_addr(&self, dst_addr: SocketAddrV4) -> std::io::Result<SocketAddrV4> {
        let udp_sock = std::net::UdpSocket::bind("0.0.0.0:0")?;
        udp_sock.connect(dst_addr)?;
        let addr = udp_sock.local_addr()?;
        match addr {
            std::net::SocketAddr::V4(v4) => Ok(v4),
            std::net::SocketAddr::V6(_) => Err(std::io::Error::new(
                std::io::ErrorKind::Other,
                "ipv6 local addr not supported",
            )),
        }
    }
    pub fn send_pkt(&self, src_addr: SocketAddrV4, dst_addr: SocketAddrV4, payload: &[u8]) {
        let mut packet = [0u8; 1500];

        let packet_len = 8 + payload.len();
        packet[8..packet_len].copy_from_slice(payload);

        {
            let mut udp_header = MutableUdpPacket::new(&mut packet[..packet_len]).unwrap();
            udp_header.set_source(src_addr.port());
            udp_header.set_destination(dst_addr.port());
            udp_header.set_length(packet_len as _);
            let n = udp_header.to_immutable().packet().len();

            let checksum =
                ipv4_checksum(&udp_header.to_immutable(), &src_addr.ip(), &dst_addr.ip());
            log::info!("n={}, checksum={:04x}", n, checksum);
            udp_header.set_checksum(checksum);
        }

        self.sock
            .send_to(&packet[..packet_len], &dst_addr.into())
            .expect("packet send error");
    }
}
